#include "include/INSE.h"

using namespace std;

#define PI 3.141592653589793
#define EXP 2.718281828

#define Dim 2

Real T0 = 0.5;
int times = 3;

// 记录误差
Vector err(times);
Vector errp(times);
int ind = 0;

Real Cr = 0.75;
Real Re = 300;

void test(int N, fstream &fp, fstream &fmat)
{
    // 对流常数
    Real nu = 1.0 / Re;

    // 空间和时间步长
    Real h = 1.0 / N;
    Real k = Cr * h;

    int SN = N;
    int TN = T0 / k;
    Real T = k * TN;

    // 速度场
    FuncXt u = [=](Vector v, Real t) -> Vector
    {
        Real x = v[0];
        Real y = v[1];
        Vector res(2);
        res[0] = 1 + 2 * pow(EXP, -8 * PI * PI * nu * t) * (-cos(2 * PI * (x - t)) * sin(2 * PI * (y - t)));
        res[1] = 1 + 2 * pow(EXP, -8 * PI * PI * nu * t) * (sin(2 * PI * (x - t)) * cos(2 * PI * (y - t)));

        return res;
    };

    // 压力场
    FuncXt p = [=](Vector v, Real t) -> Vector
    {
        Real x = v[0];
        Real y = v[1];
        Vector res(1);
        res[0] = -pow(EXP, -16 * PI * PI * nu * t) * (cos(4 * PI * (x - t)) + cos(4 * PI * (y - t)));

        return res;
    };

    // 创建单位区域
    Domain<Dim> domain;

    // 积分算子
    Integrator<3, Dim> I;

    // 包装初值
    FuncX U = [=](Vector x)
    {
        return u(x, 0);
    };

    // 初始化网格
    GridX<Dim> U0(Dim, Grid<Dim>(SN + 1));

    // 填充初值
    I.fill(domain, U0, U);

    // 求解器
    INSESolver<Dim> insesolver(p, k, nu);

    // 创建压力容器
    Grid<Dim> P(SN + 1, 0);

    // 开始计算
    Real t = insesolver.solve(U0, P, TN);

    // 包装终值
    FuncX UE = [=](Vector x)
    {
        return u(x, T);
    };
    FuncX PE = [=](Vector x)
    {
        return p(x, T);
    };

    // 初始化网格
    GridX<Dim> U1(Dim, Grid<Dim>(SN + 1));
    Grid<Dim> P1(SN + 1);

    // 填充终值
    I.fill(domain, U1, UE);
    I.fill(domain, P1, PE);

    // 估计误差
    Real err1 = TensorOp::grid(U1[0] - U0[0]);
    Real err2 = TensorOp::grid(U1[1] - U0[1]);

    fp << "SN: " << SN << endl;
    fp << "Time: " << t << endl;

    // 记录网格误差
    err[ind] = (err1 + err2) / 2;
    errp[ind] = TensorOp::grid(P - P1);
    fp << "Error of u: " << err[ind] << endl;
    fp << "Error of p: " << errp[ind] << endl;
    fp << endl;
    ind++;

    // 输出绘图程序
    fmat << "XY = [" << endl;

    // 循环输出坐标
    for (int j = 0; j < N + 1; j++)
    {
        Real h = 1.0 / (N + 1);
        fmat << h * j << " " << h * j << endl;
    }

    fmat << "];" << endl;

    // 绘制速度
    fmat << "[X, Y] = meshgrid(XY(:,1), XY(:,2));" << endl;
    fmat << "U = [" << U0[0] << "];" << endl;
    fmat << "figure;" << endl;
    fmat << "plot3(X,Y,U,'.b','MarkerSize',5);" << endl;
    fmat << "title('N = " << N << ", ux');" << endl;

    fmat << "U = [" << U0[1] << "];" << endl;
    fmat << "figure;" << endl;
    fmat << "plot3(X,Y,U,'.b','MarkerSize',5);" << endl;
    fmat << "title('N = " << N << ", uy');" << endl;

    // 绘制压力
    fmat << "P = [" << P << "];" << endl;
    fmat << "figure;" << endl;
    fmat << "plot3(X,Y,P,'.b','MarkerSize',5);" << endl;
    fmat << "title('N = " << N << ", p');" << endl;
}

int main()
{
    int N = 16;

    fstream fp("test32", ios::out);
    fstream fmat("test32.m", ios::out);

    // 记录每次的解
    GridX<Dim> U(times);

    // 测试并记录误差
    for (int i = 0; i < times; i++)
    {
        test(N, fp, fmat);
        N *= 2;
    }

    // 计算收敛阶
    fp << "Convergence rate of u: " << TensorOp::rate(err) << endl;
    fp << "Convergence rate of p: " << TensorOp::rate(errp) << endl;

    fp << endl;

    fp.close();
    fmat.close();

    return 0;
}